/*
 * lmf_main.cpp
 *
 * Copyright 2013 Hua Ouyang (houyang@gatech.edu)
 *
 * main file for Data-distributed low-rank maxtrix factorization program
 * Example run (use --help for help on parameters)
 *  Generate data
 *   mpirun -np 16 ./lmf  --gen --data=m1000_10_200_16 --m=1000 --n=200 --k=10 
 *  Decentralized Low Rank Matrix Factorization
 *   mpirun -np 2 ./lmf --lf-dist --k=10 --data=m200_10_200_2 --alpha 1 --rho 10 --graph equal --degree 1 --iter=2000 --interval=200 --iterinfo=log_200_10_200_np2_d1
 *   mpirun -np 16 ./lmf --lf-dist --k=10 --data=m1000_10_200_16 --alpha 1 --rho 10 --graph equal --degree 2 --iter=2000 --interval=200 --iterinfo=log_1000_10_200_np16_d1
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#define ARMA_USE_LAPACK

#include <iostream>
#include <armadillo>

#include "helper/helper.h"
#include "algorithm/StopWatch.h"
#include <boost/mpi.hpp>
boost::mpi::communicator* p_world = NULL; // p_world is intentionaly placed before
#include "mpi_helper.h"                   // including mpi_helper.h

#include <opencv2/opencv.hpp>
#include <iomanip>

po::variables_map vm;
po::options_description desc("Allowed options");

// JOBS
void run_anls();
void run_master_slave_anls();
void run_dist_anls();
void run_dist_lmf();
void generate_data();

int main(int argc, char** argv)
{
  mpi::environment env(argc, argv);
  mpi::communicator world;
  p_world = &world;
  
  //std::cerr << "rank=" << MPI_RANK << std::endl;

  desc.add_options()
    ("help", "produce help message")
    ("verbose", "print all debug information")
    ("interval", po::value<int>()->default_value(100), "data dimension hint")
    // DATA
    ("data", po::value<std::string>(), "data file")
    ("output", po::value<std::string>(), "output basis file")
    ("iterinfo", po::value<std::string>(), "iteration information file")
    // MODEL
    ("model", po::value<std::string>(), "output model file")
    // METHODS
    ("gen", "generate random data (MPI_SIZE)\n--m --n --k")
    ("anls-seq", "solve A = WH^T by Sequential Alternating Nonnegative Least Square\n--data --rate --rho --iter --k")
    ("anls-dist", "solve A = WH^T by Distributed Alternating Nonnegative Least Square\n--data --rate --rho --graph [--degree] --iter --k")
    ("lf-dist", "solve A = UV by Distributed Low-Rank Matrix Factorization\n--data --alpha --rho --graph [--degree] --iter --k")
    ("image", "read images")
    // COMMON PARAMETERS
    ("iter", po::value<int>()->default_value(1000), "maximum iterations")
    ("rate", po::value<double>()->default_value(1.0), "learning rate / step size")
    ("density", po::value<double>()->default_value(0.3), "sparsity density")
    ("rho", po::value<double>()->default_value(0.1), "augmented penalty")
    ("alpha", po::value<double>()->default_value(0.1), "augmented penalty for Low-Rank MF")
    ("normalize", "keep columns of H normalized (norm <= 1)")
    // DISTRIBUTED ALGORITHMS
    ("graph", po::value<std::string>(), "decentralized ADMM with bipartie graph (grid|equal)")
    ("degree", po::value<int>(), "decentralized ADMM with bipartie graph")
    // GENERATE DATA
    ("m", po::value<int>()->default_value(10), "= A.nrows")
    ("n", po::value<int>()->default_value(10), "= A.ncols")
    ("k", po::value<int>()->default_value(2), "= W.ncols = H.nrows")
    ("np", po::value<int>()->default_value(2), "Number of processes to generate data")
    ;

  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm); 

  if (vm.count("help")) {
    std::cerr << desc << std::endl;
  }
  else if (vm.count("anls-seq")) {
    MSG("Sequential ANLS");
    run_anls();
  }
  else if (vm.count("anls-dist")) {
    if (vm.count("graph")) { // graph --> Decentralized ALS
        run_dist_anls();
      }
    else
      run_master_slave_anls(); // no graph --> Master/Slave ALS
  }
  else if (vm.count("lf-dist")) {
    if (vm.count("graph")) {
      run_dist_lmf();  // graph --> Decentralized Low-Rank MF
    }
    else {
      //run_master_slave_lf(); // no graph --> Master/Slave Low-Rank MF
    }
  }
  else if (vm.count("gen")) { // generate data
    generate_data();
  }
  else if (vm.count("image")) {
    PO_VARIABLE(filename, vm, "data", std::string);
    int m = 2429;
    int n = 19*19;    
    arma::mat A(m, n);
    for (int r = 1; r <= 2429; r++) {
      std::stringstream ss;
      ss << filename << std::setw(4) << std::setfill('0') << r << ".pgm";
      MSG_VERBOSE(vm,ss.str());
      cv::Mat img = cv::imread(ss.str(), 1), gray;
      //MSG(img.rows << " " << img.cols << " " << img.type()
      //  << " " << img.depth() << " " << img.channels());
      cv::cvtColor(img, gray, CV_BGR2GRAY);
      //MSG(gray.rows << " " << gray.cols << " " << gray.type()
      //  << " " << gray.depth() << " " << gray.channels() << " " << CV_8UC1);
      cv::Mat_<unsigned char>& g = (cv::Mat_<unsigned char>&) gray;
      int k=0;
      for(int j = 0; j < g.cols; j++)
        for(int i = 0; i < g.rows; i++,k++) A(r-1,k)=((double)g(i,j)) / 255;
    }
    A.save(filename, arma::arma_binary);
    MSG_VERBOSE(vm,A);
    
    PO_VARIABLE(np, vm, "np", int);
    int start = 0;
    for (int i = 0; i < np; i++, start += m/np) {
      arma::mat B = A.rows( start, i < np-1 ? start + (m/np) - 1 : m-1 );
      std::stringstream ss;
      ss << filename << "_" << np << "_" << i;
      B.save(ss.str(), arma::arma_binary);
    }
  }
  else {
    MSG(desc);
  }

  return 0;
}

void generate_data()
{
  PO_REQUIRE(vm, "data");

  PO_VARIABLE(m, vm, "m", int);
  PO_VARIABLE(n, vm, "n", int);
  PO_VARIABLE(k, vm, "k", int);
  PO_VARIABLE(data, vm, "data", std::string);
  
  using namespace arma;
  srand(1+MPI_RANK);

  mat W = randn<mat>(m/MPI_SIZE, k);
  mat H;
  if (MPI_RANK == 0) H = randn<mat>(n, k);
  mpi::broadcast(*p_world, H, 0);

  mat A = W*trans(H);
  std::stringstream ss;
  ss << data << "_" << MPI_RANK;
  A.save(ss.str(), arma_binary);
}
